package top.kangyaocoding.tech.controller;

import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

/**
 * 描述: AI 聊天
 *
 * @author K·Herbert
 * @since 2025-04-13 20:05
 */
@Slf4j
@RestController
@RequestMapping("/api/v1/ai")
@CrossOrigin(origins = "*")
public class ChatController {

    private static final String DEFAULT_PROMPT = "你好，你是谁？";
    private final ChatModel chatModel;

    public ChatController(ChatModel chatModel) {
        this.chatModel = chatModel;
    }

    /**
     * 简单调用。一次性输出全部结果
     *
     * @return String types.
     */
    @GetMapping("/simple/chat")
    public String simpleChat(HttpServletResponse response) {
        // 避免返回乱码
        response.setCharacterEncoding("UTF-8");

        return chatModel.call(
                new Prompt(DEFAULT_PROMPT)
        ).getResult().getOutput().getContent();
    }

    /**
     * Stream 流式调用。输出的结果实现打字机效果。
     * 注意要配置 MediaType.TEXT_EVENT_STREAM_VALUE,
     * 这样配置说明定义了一个请求的 Header,
     * 并设置为 Content-Type：text/event-stream;charset=UTF-8,
     * 这样前端才能识别为流式请求。
     * @return Flux<String> types.
     */
    @PostMapping(name = "/stream/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamChat(HttpServletResponse response) {

        // 避免返回乱码
        response.setCharacterEncoding("UTF-8");

        Flux<ChatResponse> stream = chatModel.stream(
                new Prompt(
                        DEFAULT_PROMPT
                )
        );
        return stream.map(result -> result.getResult().getOutput().getContent());
    }
}
